package com.kehao.newcoder;

import org.junit.Test;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.PriorityQueue;

public class NC51Solution {

    private class ListNode {
        int val;
        ListNode next;

        ListNode(int x) {
            val = x;
            next = null;
        }
    }

    public ListNode mergeKLists(ArrayList<ListNode> lists) {
        ListNode head = new ListNode(0);
        ListNode p = head;
        lists.removeIf(i->i==null);

        PriorityQueue<ListNode> priorityQueue = new PriorityQueue<>(Comparator.comparingInt(node -> node.val));
        priorityQueue.addAll(lists);

        while (priorityQueue!=null&&!priorityQueue.isEmpty()){
            ListNode minListNode = priorityQueue.poll();
            ListNode tmp = minListNode;
            minListNode = minListNode.next;
            p.next = tmp;
            p = tmp;
            if(minListNode!=null){
                priorityQueue.add(minListNode);
            }
        }
        return head.next;
    }

    @Test
    public void test(){
        ListNode node1 = generateList(new int[]{1,2,3});
        ListNode node2 = generateList(new int[]{4,5,6,7});
        ArrayList<ListNode> lists = new ArrayList<>();
        lists.add(node1);
        lists.add(node2);
        ListNode listNode = mergeKLists(lists);
        printList(listNode);
    }

    private void printList(ListNode listNode){
        while (listNode!=null){
            System.out.print(listNode.val+" ");
            listNode=listNode.next;
        }
    }

    private ListNode generateList(int[] arr) {
        ListNode head = new ListNode(0);
        ListNode p = head;
        for (int i = 0; i < arr.length; i++) {
            ListNode listNode = new ListNode(arr[i]);
            p.next = listNode;
            p = p.next;
        }
        return head.next;
    }
}
